#include"hfsolver_k.h"
// void SHFSolver_WS::init()
// {
//   //initialize HF Orbitals with hamiltonian basis
//   int totalOrbitals=pSystem_Tz->Orbitals.size();
//   HFOrbitals.resize(totalOrbitals);
//   for(int i=0;i<pSystem_Tz->Groups.size();i++)
//     {
//       int numOrbInAGroup=pSystem_Tz->Groups[i].size();
//       for(int j=0;j<numOrbInAGroup;j++)
// 	{
// 	  int orb=pSystem_Tz->Groups[i][j];
// 	  const OrbitalType & Orb=pSystem_Tz->Orbitals[orb];
// 	  HFOrbitals[orb].set(i,Orb.l,Orb.jj,Orb.tz,Orb.e);
// 	  HFOrbitals[orb].Coeff.resize(numOrbInAGroup);
// 	  //clear all coeff
// 	  for(int indice=0;indice<numOrbInAGroup;indice++)
// 	    {
// 	      HFOrbitals[orb].Coeff[indice]=0;
// 	    }
// 	  //set initial coeff
// 	  HFOrbitals[orb].Coeff[j]=DataType(1);
// 	}
//     }
//   fillOrbitals();
// }

// void SHFSolver_WS::fillOrbitals()
// {
//   stable_sort( HFOrbitals.begin(),HFOrbitals.end() );
//   int resZ=Z;
//   int resN=A-Z;
//   for(int i=0;i<HFOrbitals.size();i++)
//     {
//       //fill proton
//       if(HFOrbitals[i].tz==-1)
// 	{
// 	  int num=HFOrbitals[i].jj+1;
// 	  resZ-=num;
// 	  if(resZ >= 0)
// 	    {
// 	      HFOrbitals[i].state=0;//hole
// 	      HFOrbitals[i].OccuP=1.0;
// 	    }
// 	  else if( (resZ + num) > 0)
// 	    {
// 	      HFOrbitals[i].state=0;//hole, but not fully occupied.
// 	      HFOrbitals[i].OccuP=double(resZ + num)/num;
// 	    }
// 	  else
// 	    {
// 	      HFOrbitals[i].state=1;//particle
// 	      HFOrbitals[i].OccuP=0.0;
// 	    }
// 	}
//       //fill neutron
//       if(HFOrbitals[i].tz==1)
// 	{
// 	  int num=HFOrbitals[i].jj+1;
// 	  resN-=num;
// 	  if(resN >= 0)
// 	    {
// 	      HFOrbitals[i].state=0;//hole
// 	      HFOrbitals[i].OccuP=1.0;
// 	    }
// 	  else if( (resN + num) > 0)
// 	    {
// 	      HFOrbitals[i].state=0;//hole, but not fully occupied.
// 	      HFOrbitals[i].OccuP=double(resN + num)/num;
// 	    }
// 	  else
// 	    {
// 	      HFOrbitals[i].state=1;//particle
// 	      HFOrbitals[i].OccuP=0.0;
// 	    }

// 	}
//     }
// }


// void SHFSolver_WS::diag()
// {
//   int totalOrbitals=HFOrbitals.size();
//   vector< SHF_Orbital<DataType> > newHFOrbitals(totalOrbitals);
  
//   int num=0;
//   int totalGroups=pSystem_Tz->Groups.size();
//   for(int i=0;i<totalGroups;i++)
//     {
//       int orb=pSystem_Tz->Groups[i][0];
//       const OrbitalType & Orb=pSystem_Tz->Orbitals[orb];
//       int l=Orb.l;
//       int jj=Orb.jj;
//       int tz=Orb.tz;
      
//       int dim=pSystem_Tz->Groups[i].size();
//       Mat mat(dim,dim);
//       for(int a=0;a<dim;a++)
// 	{
// 	  for(int b=a;b<dim;b++)
// 	    {
// 	      mat(a,b)=HFhamil(i,a,b);
// 	      mat(b,a)=mat(a,b);
// 	    }
// 	}
//       /* cout<<"==============\n"; */
//       /* cout<<mat<<endl; */
//       /* cout<<"==============\n"; */
//       //     SelfAdjointEigenSolver<Mat> es(mat);
//       ComplexEigenSolver<Mat> es(mat);
//       Vec eigenVals=es.eigenvalues();
//       Mat eigenVecs=es.eigenvectors();
//       for(int j=0;j<dim;j++)
// 	{
// 	  newHFOrbitals[num+j].set(i,l,jj,tz,eigenVals[j]);
// 	  newHFOrbitals[num+j].Coeff.resize(dim);
// 	  //Coeff
// 	  complexd norm=sqrt( eigenVecs.col(j).transpose()*eigenVecs.col(j) );
// 	  for(int k=0;k<dim;k++)
// 	    newHFOrbitals[num+j].Coeff[k]=eigenVecs(k,j)/norm;
// 	}
//       num+=dim;
//     }
//   HFOrbitals=newHFOrbitals;
//   fillOrbitals();
// }

// void SHFSolver_WS::iter()
// {
//   double sigma=0;
//   double eps=1e-6;
//   const int maxIter=100;
//   init();
//   int totalOrbitals=HFOrbitals.size();
//   int iter=0;
//   cout<<"begin  HF. iteration\n";
//   do
//     {     
//       vector< SHF_Orbital<DataType> > oldHFOrbitals=HFOrbitals;
//       //      printOrbitals();
//       diag();
//       sigma=0;
//       for(int i=0;i<totalOrbitals;i++)
//   	{
//   	  sigma+=abs(oldHFOrbitals[i].e-HFOrbitals[i].e);
//   	}
//       sigma/=totalOrbitals;
//       ++iter;
//       cout<<iter<<"\t"<<sigma<<endl;
//     }
//   while(sigma>eps && iter<maxIter);
//   cout<<"end  HF. iteration\n";
// }

// void SHFSolver_WS::calEnergy()
// {
//   Energy=0;
//   for(int i=0;i<HFOrbitals.size();i++)
//     {
//       if(HFOrbitals[i].state) continue;//skip particles
//       int dim=HFOrbitals[i].Coeff.size();
//       int group=HFOrbitals[i].GroupIndex;
//       double Possi=HFOrbitals[i].OccuP;
//       int jj1=HFOrbitals[i].jj+1;
//       for(int a=0;a<dim;a++)
// 	{
// 	  for(int b=0;b<dim;b++)
// 	    {
// 	      int bra=pSystem_Tz->Groups[group][a];
// 	      int ket=pSystem_Tz->Groups[group][b];	      
// 	      Energy+=( getOneBody(bra,ket)+0.5*SHFDiagram(bra,ket) )*Possi*HFOrbitals[i].Coeff[a]*HFOrbitals[i].Coeff[b]*DataType(jj1);
// 	    }
// 	}
//     }
// }

//============================================================================


void SHFSolver_K::init()
{
  double hbar_omega=15;
  double b=hbar_c/sqrt(mc2*hbar_omega);
  //initialize HF Orbitals
  int totalOrbitals=pSystem_Tz->Orbitals.size();
  HFOrbitals.resize(totalOrbitals);
  for(int i=0;i<pSystem_Tz->Groups.size();i++)
    {
      int numOrbInAGroup=pSystem_Tz->Groups[i].size();
      for(int n=0;n<numOrbInAGroup;n++)
	{
	  int orb=pSystem_Tz->Groups[i][n];
	  const OrbitalType & Orb=pSystem_Tz->Orbitals[orb];
	  HFOrbitals[orb].set(i,Orb.l,Orb.jj,Orb.tz,2*n+Orb.l);
	  HFOrbitals[orb].Coeff.resize(numOrbInAGroup);
	  //set coeff
	  for(int indice=0;indice<numOrbInAGroup;indice++)
	    {
	      double k=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[i][indice] ].k;
	      double wk=pSystem_Tz->Orbitals[ pSystem_Tz->Groups[i][indice] ].wk;
	      HFOrbitals[orb].Coeff[indice]=ho_k(n,Orb.l,b,k)*k*sqrt(wk);
	    }
	}
    }
  fillOrbitals();
}

SHFSolver_K::DataType SHFSolver_K::SHFDiagram(int group,int a,int c)
  {
    int bra=pSystem_Tz->Groups[group][a];
    int ket=pSystem_Tz->Groups[group][c];

    int jja=pSystem_Tz->Orbitals[bra].jj;
    int tza=pSystem_Tz->Orbitals[bra].tz;
    int la=pSystem_Tz->Orbitals[bra].l;
    double ka=pSystem_Tz->Orbitals[bra].k;

    double kc=pSystem_Tz->Orbitals[ket].k;

    DataType val2B=0;
    DataType valk1k2=0;
    for(int i=0;i<HFOrbitals.size();i++)
      {
	if(HFOrbitals[i].state) continue;//skip particles
	int jjh=HFOrbitals[i].jj;
	int Jmin=abs(jja - jjh)/2;
	int Jmax=(jja + jjh)/2;

	if(Jmin>Jmax) continue;

	int GroupIndex=HFOrbitals[i].GroupIndex;
	int dim=pSystem_Tz->Groups[GroupIndex].size();

	for(int k=0;k<dim;k++)
	  for(int l=0;l<dim;l++)
	    {
	      for(int J=Jmin;J<=Jmax;J++)
		{
		  int brak=pSystem_Tz->Groups[GroupIndex][k];
		  int ketl=pSystem_Tz->Groups[GroupIndex][l];
		  val2B += getTwoBody(bra,brak,ket,ketl,J) * HFOrbitals[i].Coeff[k] * HFOrbitals[i].Coeff[l] * HFOrbitals[i].OccuP * DataType(2*J+1);
		}
	    }
	//k1 \dot k2
	if(HFOrbitals[i].tz!=tza) continue;
	if((HFOrbitals[i].l+la+1)%2) continue;//parity
	for(int J=Jmin;J<=Jmax;J++)
	  {
	    valk1k2-=(2*J+1)*(jjh+1) * pow(gsl_sf_coupling_3j(jja,2,jjh,1,0,-1),2) * gsl_sf_coupling_6j(jja,jjh,J*2,jja,jjh,2) * ka * kc * HFOrbitals[i].Coeff[a] * HFOrbitals[i].Coeff[c];
	  }

      }
    val2B/=(jja+1);
    return val2B-valk1k2*hbar_c*hbar_c/(A*mc2);
  }



void SHFSolver_K::fillOrbitals()
{
  stable_sort( HFOrbitals.begin(),HFOrbitals.end() );
  int resZ=Z;
  int resN=A-Z;
  for(int i=0;i<HFOrbitals.size();i++)
    {
      //fill proton
      if(HFOrbitals[i].tz==-1)
	{
	  int num=HFOrbitals[i].jj+1;
	  resZ-=num;
	  if(resZ >= 0)
	    {
	      HFOrbitals[i].state=0;//hole
	      HFOrbitals[i].OccuP=1.0;
	    }
	  else if( (resZ + num) > 0)
	    {
	      HFOrbitals[i].state=0;//hole, but not fully occupied.
	      HFOrbitals[i].OccuP=double(resZ + num)/num;
	    }
	  else
	    {
	      HFOrbitals[i].state=1;//particle
	      HFOrbitals[i].OccuP=0.0;
	    }
	}
      //fill neutron
      if(HFOrbitals[i].tz==1)
	{
	  int num=HFOrbitals[i].jj+1;
	  resN-=num;
	  if(resN >= 0)
	    {
	      HFOrbitals[i].state=0;//hole
	      HFOrbitals[i].OccuP=1.0;
	    }
	  else if( (resN + num) > 0)
	    {
	      HFOrbitals[i].state=0;//hole, but not fully occupied.
	      HFOrbitals[i].OccuP=double(resN + num)/num;
	    }
	  else
	    {
	      HFOrbitals[i].state=1;//particle
	      HFOrbitals[i].OccuP=0.0;
	    }

	}
    }
}


void SHFSolver_K::diag()
{
  int totalOrbitals=HFOrbitals.size();
  vector< SHF_Orbital<DataType> > newHFOrbitals(totalOrbitals);
  
  int num=0;
  int totalGroups=pSystem_Tz->Groups.size();
  for(int i=0;i<totalGroups;i++)
    {
      int orb=pSystem_Tz->Groups[i][0];
      const OrbitalType & Orb=pSystem_Tz->Orbitals[orb];
      int l=Orb.l;
      int jj=Orb.jj;
      int tz=Orb.tz;
      
      int dim=pSystem_Tz->Groups[i].size();
      Mat mat(dim,dim);
      for(int a=0;a<dim;a++)
	{
	  for(int b=a;b<dim;b++)
	    {
	      mat(a,b)=HFhamil(i,a,b);
	      mat(b,a)=mat(a,b);
	    }
	}
      // cout<<"==============\n";
      // cout<<mat<<endl;
      // cout<<"==============\n";
      SelfAdjointEigenSolver<Mat> es(mat);
      Vec eigenVals=es.eigenvalues();
      Mat eigenVecs=es.eigenvectors();
      for(int j=0;j<dim;j++)
	{
	  newHFOrbitals[num+j].set(i,l,jj,tz,eigenVals[j]);
	  newHFOrbitals[num+j].Coeff.resize(dim);
	  //Coeff
	  for(int k=0;k<dim;k++)
	    newHFOrbitals[num+j].Coeff[k]=eigenVecs(k,j);
	}
      num+=dim;
    }
  HFOrbitals=newHFOrbitals;
  fillOrbitals();
}

void SHFSolver_K::iter()
{
  double sigma=0;
  double eps=1e-4;
  const int maxIter=100;
  init();
  int totalOrbitals=HFOrbitals.size();
  numIter=0;
  TwoBodyMat.clear();
  do
    {
      count=0;
      vector< SHF_Orbital<DataType> > oldHFOrbitals=HFOrbitals;
      printOrbitals(10);
      diag();
      sigma=0;
      for(int i=0;i<totalOrbitals;i++)
  	{
  	  sigma+=abs(oldHFOrbitals[i].e-HFOrbitals[i].e);
  	}
      sigma/=totalOrbitals;
      ++numIter;
      cout<<sigma<<endl;
    }
  while(sigma>eps && numIter<maxIter);
  numIter=0;
}

void SHFSolver_K::calEnergy()
{
  Energy=0;
  for(int i=0;i<HFOrbitals.size();i++)
    {
      if(HFOrbitals[i].state) continue;//skip particles
      int dim=HFOrbitals[i].Coeff.size();
      int group=HFOrbitals[i].GroupIndex;
      double Possi=HFOrbitals[i].OccuP;
      int jj1=HFOrbitals[i].jj+1;

      for(int a=0;a<dim;a++)
	{
	  for(int b=0;b<dim;b++)
	    {
	      int bra=pSystem_Tz->Groups[group][a];
	      int ket=pSystem_Tz->Groups[group][b];
	      Energy+=( getOneBody(bra,ket)+0.5*SHFDiagram(group,a,b) )*Possi*HFOrbitals[i].Coeff[a]*HFOrbitals[i].Coeff[b]*DataType(jj1);
	    }
	}
    }
}
